package org.loong.crypto.extension.cms.jce;

import java.io.IOException;
import java.security.PublicKey;
import java.security.cert.X509Certificate;
import java.security.interfaces.ECPublicKey;
import java.security.interfaces.RSAPublicKey;

import org.bouncycastle.asn1.ASN1Encoding;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.asn1.x9.X9ObjectIdentifiers;
import org.bouncycastle.operator.AsymmetricKeyWrapper;
import org.bouncycastle.operator.GenericKey;
import org.bouncycastle.operator.OperatorException;
import org.loong.crypto.common.exception.CryptoException;
import org.loong.crypto.core.algorithm.CipherAlgorithm;
import org.loong.crypto.core.algorithm.EncryptionMethod;
import org.loong.crypto.core.params.CipherParameters;
import org.loong.crypto.extension.operator.OperatorUtils;
import org.loong.crypto.extension.operator.SM2Utils;
import org.loong.crypto.service.CryptoService;

public class JceAsymmetricKeyWrapper extends AsymmetricKeyWrapper {

    private CryptoService cryptoService;

    private PublicKey publicKey;

    public JceAsymmetricKeyWrapper(CryptoService cryptoService, X509Certificate certificate) {
        this(cryptoService, certificate.getPublicKey());
    }

    public JceAsymmetricKeyWrapper(CryptoService cryptoService, PublicKey publicKey) {
        super(SubjectPublicKeyInfo.getInstance(publicKey.getEncoded()).getAlgorithm());

        this.cryptoService = cryptoService;
        this.publicKey = publicKey;
    }

    /**
     * Generate wrapped key.
     * 
     * @param encryptionKey the encryptionKey
     * @return the key bytes
     * @throws OperatorException {@link OperatorException}
     */
    public byte[] generateWrappedKey(GenericKey encryptionKey) throws OperatorException {
        try {
            byte[] keyBytes = OperatorUtils.getKeyBytes(encryptionKey);
            CipherAlgorithm algorithm = null;
            CipherParameters.Builder paramsBuilder = CipherParameters.builder().key(publicKey);
            ASN1ObjectIdentifier algorithmId = getAlgorithmIdentifier().getAlgorithm();
            if (algorithmId.equals(PKCSObjectIdentifiers.rsaEncryption)) {
                algorithm = CipherAlgorithm.RSA;
                paramsBuilder.encryptionMethod(EncryptionMethod.RSA_ECB_PKCS1);
            } else if (algorithmId.equals(X9ObjectIdentifiers.id_ecPublicKey)) {
                algorithm = CipherAlgorithm.SM2;
            } else {
                throw new OperatorException("Unsupported algorithm: " + algorithmId);
            }

            byte[] encryptedKeyBytes = cryptoService.encrypt(algorithm, paramsBuilder.build(), keyBytes);
            if (encryptedKeyBytes == null) {
                throw new OperatorException("Unable to encrypt contents key.");
            }
            
            if (publicKey instanceof RSAPublicKey) {
                return encryptedKeyBytes;
            } else if (publicKey instanceof ECPublicKey) {
                return SM2Utils.convertSM2Cipher(encryptedKeyBytes).getEncoded(ASN1Encoding.DER);
            } else {
                throw new OperatorException("Unknown public key type.");
            }
        } catch (CryptoException | IOException e) {
            throw new OperatorException(e.getMessage(), e);
        }
    }
}
